﻿#include <ntddk.h>

#include "driver.h"


#ifdef NDEBUG
#define LOG(...)
#else
#define LOG(...) DbgPrintEx(DPFLTR_IHVDRIVER_ID, DPFLTR_INFO_LEVEL, __VA_ARGS__)
#endif

NTSTATUS DriverCreateDispatch(
    IN struct _DEVICE_OBJECT* DeviceObject,
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    PAGED_CODE();

    UNREFERENCED_PARAMETER(DeviceObject);
    LOG("[Debug] DriverCreateDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverCreateDispatch Leave\r\n");

    return status;
}

NTSTATUS 
DriverCloseDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject,
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    PAGED_CODE();

    UNREFERENCED_PARAMETER(DeviceObject);
    LOG("[Debug] DriverCloseDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverCloseDispatch Leave\r\n");

    return status;
}

NTSTATUS 
DriverReadDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject, 
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    PAGED_CODE();

    UNREFERENCED_PARAMETER(DeviceObject);
    LOG("[Debug] DriverReadDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverReadDispatch Leave\r\n");

    return status;
}

NTSTATUS 
DriverWriteDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject, 
    _Inout_ struct _IRP* Irp
    )
{
    NTSTATUS status;

    PAGED_CODE();

    UNREFERENCED_PARAMETER(DeviceObject);
    LOG("[Debug] DriverWriteDispatch Enter\r\n");
    status = STATUS_SUCCESS;
    Irp->IoStatus.Status = status;
    Irp->IoStatus.Information = 0;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    LOG("[Debug] DriverWriteDispatch Leave\r\n");

    return status;
}

NTSTATUS
DefaultPnpHandler(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    )
{
    PAGED_CODE();
    
    LOG("[Debug] Enter DefaultPnpHandler\n");
    // 本层设备不处理，将 IRP 传给下一层设备处理
    IoSkipCurrentIrpStackLocation(Irp);
    LOG("[Debug] Leave DefaultPnpHandler\n");
    return IoCallDriver(DeviceObject, Irp);
}

NTSTATUS
RemoveDevice(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
)
{
    NTSTATUS status;
    PDEVICE_EXTENSION deviceExtension;

    PAGED_CODE();

    LOG("[Debug] Enter RemoveDevice\n");
    deviceExtension = (PDEVICE_EXTENSION)DeviceObject->DeviceExtension;
    LOG("[Debug] deviceObject=%p, Self=%p, PhysicalDeviceObject=%p, TopOfStack=%p, DeviceName=%wZ, SymbolicLinkName=%wZ\r\n",
        DeviceObject,
        deviceExtension->Self,
        deviceExtension->PhysicalDeviceObject,
        deviceExtension->TopOfStack,
        &deviceExtension->DeviceName,
        &deviceExtension->SymbolicLinkName
    );
    // 先将 IRP_MOJOR_REMOVE_DEVICE 传下去
    status = DefaultPnpHandler(DeviceObject, Irp);
    IoDeleteSymbolicLink(&deviceExtension->SymbolicLinkName);
    if (deviceExtension->PhysicalDeviceObject)
    {
        IoDetachDevice(deviceExtension->PhysicalDeviceObject);
    }
    IoDeleteDevice(DeviceObject);
    LOG("[Debug] Leave RemoveDevice\n");

    return status;
}

typedef NTSTATUS(*PnpHandlerRoutine)(
    IN PDEVICE_OBJECT DeviceObject,
    IN PIRP Irp
    );

NTSTATUS
DriverPnpDispatch(
    _In_ struct _DEVICE_OBJECT* DeviceObject,
    _Inout_ struct _IRP* Irp
)
{
    PIO_STACK_LOCATION stack;

    PAGED_CODE();

    // 设置 IRP_MN_PNP 的回调函数
    static PnpHandlerRoutine pnpHandlerRoutines[] = {
        DefaultPnpHandler,  //IRP_MN_START_DEVICE                 0x00
        DefaultPnpHandler,  //IRP_MN_QUERY_REMOVE_DEVICE          0x01
        RemoveDevice,       //IRP_MN_REMOVE_DEVICE                0x02
        DefaultPnpHandler,  //IRP_MN_CANCEL_REMOVE_DEVICE         0x03
        DefaultPnpHandler,  //IRP_MN_STOP_DEVICE                  0x04
        DefaultPnpHandler,  //IRP_MN_QUERY_STOP_DEVICE            0x05
        DefaultPnpHandler,  //IRP_MN_CANCEL_STOP_DEVICE           0x06
        DefaultPnpHandler,  //IRP_MN_QUERY_DEVICE_RELATIONS       0x07
        DefaultPnpHandler,  //IRP_MN_QUERY_INTERFACE              0x08
        DefaultPnpHandler,  //IRP_MN_QUERY_CAPABILITIES           0x09
        DefaultPnpHandler,  //IRP_MN_QUERY_RESOURCES              0x0A
        DefaultPnpHandler,  //IRP_MN_QUERY_RESOURCE_REQUIREMENTS  0x0B
        DefaultPnpHandler,  //IRP_MN_QUERY_DEVICE_TEXT            0x0C
        DefaultPnpHandler,  //IRP_MN_FILTER_RESOURCE_REQUIREMENTS 0x0D
        DefaultPnpHandler,  //
        DefaultPnpHandler,  //IRP_MN_READ_CONFIG                  0x0F
        DefaultPnpHandler,  //IRP_MN_WRITE_CONFIG                 0x10
        DefaultPnpHandler,  //IRP_MN_EJECT                        0x11
        DefaultPnpHandler,  //IRP_MN_SET_LOCK                     0x12
        DefaultPnpHandler,  //IRP_MN_QUERY_ID                     0x13
        DefaultPnpHandler,  //IRP_MN_QUERY_PNP_DEVICE_STATE       0x14
        DefaultPnpHandler,  //IRP_MN_QUERY_BUS_INFORMATION        0x15
        DefaultPnpHandler,  //IRP_MN_DEVICE_USAGE_NOTIFICATION    0x16
        DefaultPnpHandler,  //IRP_MN_SURPRISE_REMOVAL             0x17
    };
    static ULONG pnpHandlerRoutineCount = 
        sizeof(pnpHandlerRoutines) / sizeof(*pnpHandlerRoutines);

    stack = IoGetCurrentIrpStackLocation(Irp);
    if (pnpHandlerRoutineCount < stack->MajorFunction)
    {
        return DefaultPnpHandler(DeviceObject, Irp);
    }
    
    LOG("[Debug] MinorFunction=%#X, MajorFunction=%#X\n", 
        stack->MinorFunction, stack->MajorFunction);

    return pnpHandlerRoutines[stack->MajorFunction](DeviceObject, Irp);
}

VOID
DriverUnload(
    IN struct _DRIVER_OBJECT* DriverObject
)
{
    PAGED_CODE();
    
    UNREFERENCED_PARAMETER(DriverObject);
    LOG("[Debug] Enter DriverUnload\n");
    LOG("[Debug] Leave DriverUnload\n");
}

NTSTATUS
AddDevice(
    _In_ struct _DRIVER_OBJECT* DriverObject,
    _In_ struct _DEVICE_OBJECT* PhysicalDeviceObject
)
{
    NTSTATUS status;
    UNICODE_STRING deviceName;
    PDEVICE_OBJECT deviceObject;
    UNICODE_STRING symbolicLinkName;
    PDEVICE_EXTENSION deviceExtension;

    PAGED_CODE();
    // 创建设备
    RtlInitUnicodeString(&deviceName, DEVICE_NAME);
    status = IoCreateDevice(
        DriverObject,
        sizeof(DEVICE_EXTENSION),
        &deviceName,
        FILE_DEVICE_UNKNOWN,
        0,
        FALSE,
        &deviceObject
    );
    if (!NT_SUCCESS(status))
    {
        LOG("[Error] IoCreateDevice deviceName=%wZ, status=%#X\r\n", &deviceName, status);
        return status;
    }
    deviceObject->Flags |= (DO_BUFFERED_IO | DO_POWER_PAGABLE);
    deviceObject->Flags &= ~DO_DEVICE_INITIALIZING;

    // FDO 附加到 PDO
    deviceExtension = (PDEVICE_EXTENSION)deviceObject->DeviceExtension;
    RtlZeroMemory(deviceExtension, sizeof(deviceExtension));
    deviceExtension->TopOfStack = IoAttachDeviceToDeviceStack(deviceObject, PhysicalDeviceObject);
    deviceExtension->Self = deviceObject;
    deviceExtension->PhysicalDeviceObject = PhysicalDeviceObject;
    deviceExtension->DeviceName = deviceName;

    // 打印出设备堆栈
    LOG("[Debug] deviceObject=%p, TopOfStack=%p\r\n", deviceObject, deviceExtension->TopOfStack);
    LOG("[Debug] PhysicalDeviceObject=%p, PhysicalDeviceObject->AttachedDevice=%p\r\n", 
        PhysicalDeviceObject, PhysicalDeviceObject->AttachedDevice);
    ULONG deviceCount = 0;
    PDEVICE_OBJECT device = PhysicalDeviceObject;
    while (device)
    {
        LOG("[Debug] Index=%d, device=%p\r\n", deviceCount, device);
        ++deviceCount;
        device = device->AttachedDevice;
    }

    // 创建符号链接
    RtlInitUnicodeString(&symbolicLinkName, DEVICE_SYMBOLIC_LINK_NAME);
    status = IoCreateSymbolicLink(&symbolicLinkName, &deviceName);
    if (!NT_SUCCESS(status))
    {
        IoDeleteDevice(deviceObject);
        LOG("[Error] IoCreateSymbolicLink symbolicLinkName=%wZ, deviceName=%wZ status=%#X\r\n", 
            &symbolicLinkName, &deviceName, status);
        return status;
    }
    deviceExtension->SymbolicLinkName = symbolicLinkName;

    return status;
}


NTSTATUS 
DriverEntry(
    IN PDRIVER_OBJECT DriverObject,
    IN PUNICODE_STRING RegistryPath
    )
{
    LOG("[Debug] DriverObject DriverName=%wZ, HardwareDatabase=%wZ\r\n", 
        &DriverObject->DriverName, DriverObject->HardwareDatabase);
    LOG("[Debug] RegistryPath=%wZ\r\n", RegistryPath);

    Test();
    return STATUS_UNSUCCESSFUL;

    //DriverObject->DriverExtension->AddDevice = AddDevice;
    //DriverObject->DriverUnload = DriverUnload;
    //DriverObject->MajorFunction[IRP_MJ_CREATE] = DriverCreateDispatch;
    //DriverObject->MajorFunction[IRP_MJ_CLOSE] = DriverCloseDispatch;
    //DriverObject->MajorFunction[IRP_MJ_READ] = DriverReadDispatch;
    //DriverObject->MajorFunction[IRP_MJ_WRITE] = DriverWriteDispatch;
    //DriverObject->MajorFunction[IRP_MJ_PNP] = DriverPnpDispatch;

    //return STATUS_SUCCESS;
}